# -*- coding: utf-8 -*-
"""
基于FAISS的混合检索策略测试

@Author  : gaoyu
@Time    : 2025/7/22 17:05
@Function: 改进的混合检索策略，结合FAISS向量数据库和关键词匹配

本模块实现了一个高级的混合检索系统，整合了以下技术：

核心技术栈：
1. FAISS (Facebook AI Similarity Search): 高效的向量相似度搜索库
2. Ollama + DeepSeek R1: 本地部署的大语言模型和嵌入模型
3. LangChain: 提供统一的LLM和向量存储接口
4. 混合检索策略: 结合关键词匹配和语义向量搜索

系统架构：
┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
│   用户问题      │───▶│   混合检索器     │───▶│   最佳答案      │
└─────────────────┘    └──────────────────┘    └─────────────────┘
                              │
                    ┌─────────┴─────────┐
                    │                   │
            ┌───────▼────────┐  ┌──────▼──────┐
            │  关键词匹配    │  │  向量搜索   │
            │  (精确匹配)    │  │  (语义理解) │
            └────────────────┘  └─────────────┘

检索策略：
1. 关键词匹配：基于预定义关键词的精确匹配，适合处理明确的查询
2. 向量搜索：基于FAISS的语义相似度搜索，适合处理模糊或变形的查询
3. 智能融合：根据置信度动态选择最佳匹配结果

测试场景：
- 标准问题测试：验证精确匹配能力
- 模糊问题测试：验证语义理解和泛化能力
- 性能基准测试：评估检索速度和准确性

技术优势：
- 高效性：FAISS提供毫秒级的向量搜索
- 准确性：多策略融合提高匹配准确率
- 可扩展性：支持大规模知识库的构建
- 持久化：向量数据库可保存和重用
"""
import os
from difflib import SequenceMatcher
from typing import Dict, List, Optional, Tuple, Any

from langchain.prompts import PromptTemplate
from langchain.schema import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.llms.ollama import Ollama
from langchain_community.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from loguru import logger

# ==================== 模型配置和初始化 ====================

# 初始化 Ollama LLM 客户端
# DeepSeek R1 32B 是一个高性能的推理优化模型，特别适合逻辑推理任务
# 模型部署在内网服务器上，提供稳定的推理服务
ollama = Ollama(
    model="deepseek-r1:32b",           # 使用DeepSeek R1 32B模型
    base_url="http://10.48.0.81:11434"  # Ollama服务地址
)

# 初始化嵌入模型
# 使用同一个模型进行文本嵌入，确保语义表示的一致性
# 在生产环境中，可以考虑使用专门的嵌入模型如sentence-transformers
embeddings = OllamaEmbeddings(
    model="deepseek-r1:32b",           # 复用主模型进行嵌入
    base_url="http://10.48.0.81:11434"  # 保持服务地址一致
)

logger.info("模型初始化完成 - DeepSeek R1 32B")

# ==================== 测试数据集定义 ====================

# 经典逻辑推理问答数据集
# 这是一个著名的逻辑推理测试集，用于评估AI的逻辑思维和上下文理解能力
# 四个问题之间存在严密的逻辑关联，形成完整的推理链条
standard_qa = [
    {
        # 问题1：基础操作流程
        # 建立初始状态：大象被装入冰箱
        "question": "给大象装冰箱，总共分几步？",
        "answer": "三步：第一步打开冰箱门，第二步把大象装进去，第三步关上冰箱门。",
        "key": "elephant",
        "keywords": ["大象", "装冰箱", "几步"]
    },
    {
        # 问题2：复杂操作流程
        # 依赖问题1的结果，需要先移除大象再装入长颈鹿
        # 测试AI对前置条件和状态变化的理解
        "question": "给长颈鹿装冰箱，总共分几步？",
        "answer": "四步：第一步打开冰箱门，第二步把大象拿出来，第三步把长颈鹿装进去，第四步关上冰箱门。",
        "key": "giraffe",
        "keywords": ["长颈鹿", "装冰箱", "几步"]
    },
    {
        # 问题3：状态推理
        # 基于问题2的结果推断长颈鹿的位置
        # 测试AI对隐含信息的推理能力
        "question": "动物园开运动会，谁没来？",
        "answer": "长颈鹿没来，因为长颈鹿在冰箱里。",
        "key": "zoo",
        "keywords": ["动物园", "运动会", "没来", "缺席"]
    },
    {
        # 问题4：多层逻辑推理
        # 需要综合前面所有信息进行复合推理
        # 鳄鱼去运动会 → 河里没有鳄鱼 → 过河安全
        "question": "你过一条河，鳄鱼住在河里，你没被咬，为什么？",
        "answer": "因为鳄鱼去参加动物园运动会了。",
        "key": "river",
        "keywords": ["过河", "鳄鱼", "没被咬", "为什么"]
    }
]

logger.info(f"加载标准问答数据集，共 {len(standard_qa)} 条记录")


# ==================== 混合检索器实现 ====================

class HybridRetriever:
    """
    改进的混合检索策略实现
    
    结合关键词匹配和FAISS向量搜索的双重检索策略，提供更准确和鲁棒的
    信息检索能力。该实现针对小规模知识库进行了优化。
    
    核心特性：
    1. 双重检索：关键词匹配 + 向量语义搜索
    2. 智能融合：基于置信度的动态策略选择
    3. 容错机制：多级降级策略确保总能返回结果
    4. 性能优化：针对小规模数据集的快速检索
    
    检索流程：
    输入问题 → 关键词匹配 → 向量搜索 → 置信度评估 → 策略选择 → 返回答案
    
    Attributes:
        qa_data (List[Dict]): 问答数据集
        vectorstore (FAISS): FAISS向量存储实例
    """
    
    def __init__(self, qa_data: List[Dict], vectorstore: Optional[FAISS] = None):
        """
        初始化混合检索器
        
        Args:
            qa_data (List[Dict]): 问答数据集，每个元素包含question、answer、key、keywords
            vectorstore (Optional[FAISS]): FAISS向量存储实例，可选
        """
        self.qa_data = qa_data
        self.vectorstore = vectorstore
        
        logger.info(f"混合检索器初始化完成")
        logger.info(f"  - 数据集大小: {len(qa_data)} 条")
        logger.info(f"  - 向量存储: {'已启用' if vectorstore else '未启用'}")
        logger.info(f"  - 检索策略: 关键词匹配 + {'向量搜索' if vectorstore else '仅关键词'}")

    def keyword_match(self, question: str) -> Tuple[Optional[Dict], float]:
        """
        基于关键词的精确匹配策略
        
        使用关键词匹配和字符串相似度的组合算法，适合处理明确的查询。
        该方法对于包含预定义关键词的问题具有很高的准确性。
        
        算法流程：
        1. 关键词计数：统计问题中包含的预定义关键词数量
        2. 字符串相似度：使用SequenceMatcher计算文本相似度
        3. 加权融合：关键词权重70%，相似度权重30%
        4. 最优选择：返回得分最高的匹配结果
        
        Args:
            question (str): 输入的问题文本
        
        Returns:
            Tuple[Optional[Dict], float]: (最佳匹配项, 匹配得分)
                - 匹配项包含question、answer、key、keywords字段
                - 得分范围[0, 1]，越高表示匹配度越好
        """
        best_match = None
        best_score = 0.0

        for item in self.qa_data:
            # 关键词匹配计分
            keyword_score = 0
            for keyword in item['keywords']:
                if keyword in question:
                    keyword_score += 1
            
            # 标准化关键词得分（避免关键词数量差异的影响）
            normalized_keyword_score = keyword_score / len(item['keywords']) if item['keywords'] else 0

            # 计算字符串相似度（基于编辑距离）
            similarity = SequenceMatcher(None, question, item['question']).ratio()
            
            # 加权融合：关键词匹配更重要，因为它更精确
            total_score = normalized_keyword_score * 0.7 + similarity * 0.3

            if total_score > best_score:
                best_score = total_score
                best_match = item

        logger.debug(f"关键词匹配完成，最佳得分: {best_score:.3f}")
        return best_match, best_score

    def vector_match(self, question: str) -> Tuple[Optional[Dict], float]:
        """
        基于FAISS向量的语义匹配策略
        
        使用预训练的嵌入模型将问题转换为向量，然后在FAISS索引中进行
        高效的相似度搜索。这种方法能够理解语义相似性，处理问题的变形表达。
        
        技术原理：
        1. 文本嵌入：将输入问题转换为高维向量表示
        2. 向量搜索：在FAISS索引中查找最相似的向量
        3. 距离转换：将FAISS距离转换为相似度分数
        4. 元数据提取：返回匹配文档的原始问答数据
        
        FAISS优势：
        - 高效性：支持百万级向量的毫秒级搜索
        - 准确性：基于余弦相似度的精确匹配
        - 可扩展性：支持GPU加速和分布式部署
        
        Args:
            question (str): 输入的问题文本
        
        Returns:
            Tuple[Optional[Dict], float]: (最佳匹配项, 相似度分数)
                - 相似度分数范围[0, 1]，越高表示语义越相似
                - 如果向量存储未初始化或搜索失败，返回(None, 0)
        """
        if not self.vectorstore:
            logger.debug("向量存储未初始化，跳过向量匹配")
            return None, 0.0

        try:
            # 执行向量相似度搜索，返回最相似的1个结果
            results = self.vectorstore.similarity_search_with_score(question, k=1)
            
            if results:
                doc, distance = results[0]
                
                # FAISS返回的是L2距离，需要转换为相似度分数
                # 使用倒数函数：similarity = 1 / (1 + distance)
                # 距离越小，相似度越高；距离为0时，相似度为1
                similarity = 1.0 / (1.0 + distance)
                
                logger.debug(f"向量匹配完成，距离: {distance:.3f}, 相似度: {similarity:.3f}")
                return doc.metadata, similarity
            else:
                logger.debug("向量搜索未返回结果")
                return None, 0.0
                
        except Exception as e:
            logger.error(f"向量检索失败: {e}")
            return None, 0.0

    def get_best_answer(self, question: str) -> str:
        """
        智能混合检索策略
        
        结合关键词匹配和向量搜索的优势，通过多级决策机制选择最佳答案。
        该策略设计了完善的降级机制，确保在各种情况下都能提供合理的回答。
        
        决策逻辑：
        1. 高置信度关键词匹配：得分>0.5时优先使用，因为精确匹配更可靠
        2. 中等置信度向量匹配：得分>0.3时使用，处理语义相似的变形问题
        3. 低置信度关键词匹配：作为兜底策略，确保有基本的匹配结果
        4. 无匹配情况：返回友好的提示信息
        
        置信度阈值设计原理：
        - 关键词阈值0.5：确保至少有一半的关键词匹配或高文本相似度
        - 向量阈值0.3：平衡准确性和召回率，避免过于严格导致无结果
        
        Args:
            question (str): 输入的问题文本
        
        Returns:
            str: 最佳匹配的答案文本
        
        Example:
            >>> # 使用标准问答数据创建检索器
            >>> retriever = HybridRetriever(standard_qa, None)  # 仅使用关键词匹配
            >>> answer = retriever.get_best_answer("大象怎么装冰箱？")
            >>> print(answer)  # doctest: +ELLIPSIS
            把大象装进冰箱总共分为三步...
        """
        # 执行双重检索策略
        keyword_result, keyword_score = self.keyword_match(question)
        vector_result, vector_score = self.vector_match(question)

        # 记录检索结果用于分析
        logger.info(f"  关键词匹配得分: {keyword_score:.3f}")
        logger.info(f"  向量匹配得分: {vector_score:.3f}")

        # 智能决策：基于置信度选择最佳策略
        if keyword_score > 0.5:  
            # 高置信度关键词匹配：精确匹配优先
            logger.info(f"  ✓ 使用关键词匹配结果（高置信度）")
            return keyword_result['answer']
            
        elif vector_score > 0.3:  
            # 中等置信度向量匹配：语义理解补充
            logger.info(f"  ✓ 使用向量匹配结果（语义相似）")
            return vector_result['answer']
            
        elif keyword_result:  
            # 低置信度关键词匹配：兜底策略
            logger.info(f"  ⚠ 使用备选关键词匹配（低置信度）")
            return keyword_result['answer']
            
        else:
            # 无匹配情况：友好提示
            logger.warning(f"  ✗ 未找到匹配的答案")
            return "抱歉，我无法理解您的问题。请尝试重新表述或提供更多关键词。"


# 创建FAISS向量数据库（可选）
def create_faiss_db():
    try:
        logger.info("正在创建 FAISS 向量数据库...")
        documents = []
        for item in standard_qa:
            doc = Document(
                page_content=item['question'],
                metadata=item
            )
            documents.append(doc)

        vectorstore = FAISS.from_documents(documents, embeddings)
        logger.info("FAISS 数据库创建成功")
        return vectorstore
    except Exception as e:
        logger.error(f"FAISS 创建失败: {e}")
        return None


# 初始化检索器
logger.info("正在初始化混合检索器...")
vectorstore = None

# 尝试加载已有的向量数据库
if os.path.exists("./deepseek_faiss_qa_db"):
    try:
        vectorstore = FAISS.load_local(
            "./deepseek_faiss_qa_db", 
            embeddings,
            allow_dangerous_deserialization=True  # 添加这个参数
        )
        logger.info("已加载现有的 FAISS 数据库")
    except Exception as e:
        logger.error(f"加载失败: {e}，创建新的数据库")
        vectorstore = create_faiss_db()
else:
    vectorstore = create_faiss_db()

# 创建混合检索器
retriever = HybridRetriever(standard_qa, vectorstore)

# 测试标准问题
logger.info("\n=== 测试标准问题 ===")
test_questions = [
    "给大象装冰箱，总共分几步？",
    "给长颈鹿装冰箱，总共分几步？",
    "动物园开运动会，谁没来？",
    "你过一条河，鳄鱼住在河里，你没被咬，为什么？"
]

results = {}
key_mapping = ["elephant", "giraffe", "zoo", "river"]

for i, question in enumerate(test_questions):
    logger.info(f"\n问题 {i + 1}: {question}")
    answer = retriever.get_best_answer(question)
    results[key_mapping[i]] = answer
    logger.info(f"答案: {answer}")
    logger.info("-" * 50)

# 测试模糊问题
logger.info("\n=== 测试模糊匹配（修复版）===")
fuzzy_questions = [
    "大象怎么装进冰箱？",
    "长颈鹿装冰箱需要几步？",
    "运动会谁缺席了？",
    "过河为什么没被鳄鱼咬？"
]

for question in fuzzy_questions:
    logger.info(f"\n问题: {question}")
    answer = retriever.get_best_answer(question)
    logger.info(f"答案: {answer}")
    logger.info("-" * 30)

# 保存向量数据库
if vectorstore:
    try:
        vectorstore.save_local("./deepseek_faiss_qa_db")
        logger.info("\n向量数据库已保存")
    except Exception as e:
        logger.error(f"保存失败: {e}")
